#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
from atk.tasks.backends.lib_interface.acl_wrapper import AclFormat,AclTensorList

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul")
class GroupedMatmulAddGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.m = 0
        self.k = 0
        self.n = 0
        self.e = 0

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        self.m = random.randint(1,1024)
        self.k = random.randint(1,2048)
        self.n = random.randint(1,4096)
        x_shape = case_config.inputs[0][0].shape
        w_shape = case_config.inputs[1][0].shape
        group_list_shape = case_config.inputs[2].shape

        # x shape
        x_shape[0] = self.m # m
        x_shape[1] = self.k # k
        self.e = random.randint(1, min(x_shape[0], 4)) # e

        # weight shape
        w_shape[0] = self.e # e
        w_shape[1] = self.k # k
        w_shape[2] = self.n # n

        #groupList
        group_list_shape[0] = self.e # e
        case_config.inputs[2].range_values = [1, self.m]

        # dtype of x is same as weight's 
        x_dtype = case_config.inputs[0][0].dtype
        case_config.inputs[1][0].dtype = x_dtype
        if case_config.inputs[0][0].dtype == "bf16":
            case_config.inputs[6].range_values = "bf16"
        elif case_config.inputs[0][0].dtype == "fp16":
            case_config.inputs[6].range_values = "fp16"
        else:
            case_config.inputs[6].range_values = "fp32"

        return case_config
    